// Copyright  2004-2008 Backplane Ltd.
//
// This file is part of RegistryLookup
//
// RegistryLookup is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 3 of the License, or (at your
// option) any later version.
//
// RegistryLookup is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License
// along with RegistryLookup. If not, see <http://www.gnu.org/licenses/>.
//	
//
//   Purpose: Implementation of CGenericRequestBroker.
//   $Id$




#include "stdafx.h"
#include "GenericRequestBroker.h"
#include <atlstr.h>
#include <winreg.h>
#include <ReturnCodes.h>

// Construction.
CGenericRequestBroker::CGenericRequestBroker()
{
}

// Destruction.
CGenericRequestBroker::~CGenericRequestBroker()
{
}

// IRegistryRead.
STDMETHODIMP CGenericRequestBroker::ReadValueString(BSTR sKeyRoot, BSTR sKey, BSTR sValue, BSTR *psData)
{
	HRESULT hr = S_OK;

	// Check the arguments.
	if(sKeyRoot && sKey && psData) {
		// Ascertain the root key.
		HKEY hKeyRoot = 0;
		hr = GetHKeyFromString(sKeyRoot, hKeyRoot);
		if(hr == S_OK) {
			hr = AuthoriseValueRead(sKeyRoot, sKey, sValue);
			if(hr == S_OK) {
				CW2A sLocalKey = sKey;
				// Open the requested registry key.
				HKEY hKey = 0;
				if(RegOpenKeyEx(hKeyRoot, sLocalKey, 0, KEY_READ, &hKey) == ERROR_SUCCESS) {
					CW2A sLocalValue = sValue;
					// Ascertain the size of the data in the specified value.
					DWORD dwValueType = 0, dwValueSize = 0;
					if(RegQueryValueEx(hKey, sLocalValue, 0, &dwValueType, 0, &dwValueSize) == ERROR_SUCCESS) {
						// String values may not have been stored with a terminating null, so we must add space for one.
						if(dwValueType == REG_SZ || REG_MULTI_SZ || REG_EXPAND_SZ)
							dwValueSize++;
						LPBYTE pData = reinterpret_cast<LPBYTE>(malloc(dwValueSize));
						if(pData) {
							// Now read the data from the specified value.
							if(RegQueryValueEx(hKey, sLocalValue, 0, &dwValueType, pData, &dwValueSize) == ERROR_SUCCESS)
								// Convert the data to a string and return it to the caller.
								hr = GetStringFromData(pData, dwValueSize, dwValueType, psData);
							else {
								hr = RL_E_BADVALUE;
								SetupErrorInfo(hr, L"read", sValue && *sValue != L'\0' ? sValue : L"[default]", sKeyRoot, sKey);
							}
							free(pData);
						} else {
							hr = RL_E_NOMEM;
							SetupErrorInfo(hr, L"ReadValueString");
						}
					} else {
						hr = RL_E_BADVALUE;
						SetupErrorInfo(hr, L"read", sValue && *sValue != L'\0' ? sValue : L"[default]", sKeyRoot, sKey);
					}
				} else {
					hr = RL_E_OPENKEY;
					SetupErrorInfo(hr, sKeyRoot, sKey);
				}
			} /*else if(hr == S_FALSE) {
				hr = RL_E_UNAUTH;
				SetupErrorInfo(hr, L"read", sValue && *sValue != L'\0' ? sValue : L"[default]", sKeyRoot, sKey); 
			}*/
		} else if(hr == S_FALSE) {
			hr = RL_E_BADROOTKEY;
			SetupErrorInfo(hr, sKeyRoot, L"ReadValueString");
		}
	} else {
		hr = RL_E_BADARG;
		SetupErrorInfo(hr, L"GenericRequestBroker::ReadValueString");
	}

	return hr;
}

// ISupportErrorInfo.
STDMETHODIMP CGenericRequestBroker::InterfaceSupportsErrorInfo(REFIID riid)
{
	static const IID* arr[] = { &IID_IRegistryRead };

	for (int i=0; i < sizeof(arr) / sizeof(arr[0]); i++) {
		if (InlineIsEqualGUID(*arr[i],riid))
			return S_OK;
	}

	return S_FALSE;
}

// Request authorisation to read the specified registry value.
HRESULT CGenericRequestBroker::AuthoriseValueRead(BSTR sKeyRoot, BSTR sKey, BSTR sValue)
{
	HRESULT hr = S_FALSE;

	// Check the arguments.
	if(sKeyRoot && sKey) {
		// Setup the message.
		CString sValueParameter;
		if(sValue && *sValue != L'\0')
			sValueParameter.Format("the value '%S'", sValue);
		else
			sValueParameter = "the default value";
		CString sMessage;
		sMessage.Format("A process is attempting to read %s from the following registry key:\n\n%S\\%S\n\nDo you wish to authorise this operation?", sValueParameter, sKeyRoot, sKey);
		// Request authorisation from the user.
		if(MessageBox(0, sMessage, "Sidewinder RegistryLookup", MB_YESNO | MB_ICONQUESTION | MB_SETFOREGROUND) == IDYES)
			hr = S_OK;
	} else {
		hr = RL_E_BADARG;
		SetupErrorInfo(hr, L"GenericRequestBroker::AuthoriseValueRead");
	}

	return hr;
}

// Convert a string into its corresponding HKEY value.
HRESULT CGenericRequestBroker::GetHKeyFromString(const BSTR sKey, HKEY &hKey)
{
	HRESULT hr = S_OK;

	// Check the arguments.
	if(sKey) {
		// Check for known key types and return the appropriate handle.
		if(wcscmp(L"HKEY_CLASSES_ROOT", sKey) == 0)
			hKey = HKEY_CLASSES_ROOT;
		else if(wcscmp(L"HKEY_CURRENT_CONFIG", sKey) == 0)
			hKey = HKEY_CURRENT_CONFIG;
		else if(wcscmp(L"HKEY_CURRENT_USER", sKey) == 0)
			hKey = HKEY_CURRENT_USER;
		else if(wcscmp(L"HKEY_LOCAL_MACHINE", sKey) == 0)
			hKey = HKEY_LOCAL_MACHINE;
		else if(wcscmp(L"HKEY_PERFORMANCE_DATA", sKey) == 0)
			hKey = HKEY_PERFORMANCE_DATA;
		else if(wcscmp(L"HKEY_PERFORMANCE_NLSTEXT", sKey) == 0)
			hKey = HKEY_PERFORMANCE_NLSTEXT;
		else if(wcscmp(L"HKEY_PERFORMANCE_TEXT", sKey) == 0)
			hKey = HKEY_PERFORMANCE_TEXT;
		else if(wcscmp(L"HKEY_USERS", sKey) == 0)
			hKey = HKEY_USERS;
		else {
			hKey = 0;
			hr = S_FALSE;
		}
	} else {
		hr = RL_E_BADARG;
		SetupErrorInfo(hr, L"GenericRequestBroker::GetHKeyFromString");
	}

	return hr;
}

// Convert data from the registry into an appropriate string representation.
HRESULT CGenericRequestBroker::GetStringFromData(const LPBYTE pData, const DWORD dwDataLen, const DWORD dwDataType, BSTR *psData)
{
	HRESULT hr = S_OK;

	// Check the arguments.
	if(pData && psData) {
		if(dwDataType == REG_EXPAND_SZ || dwDataType == REG_SZ) {
			// Return a direct copy of the string.
			CComBSTR sData(CA2W(reinterpret_cast<const char *>(pData)));
			*psData = sData.Detach();
		} else if(dwDataType == REG_BINARY) {
			// Return a string representation of the binary value.
			wchar_t sBinary[33];
			_ltow(*reinterpret_cast<const DWORD *>(pData), sBinary, 2);
			CComBSTR sData(sBinary);
			*psData = sData.Detach();
		} else if(dwDataType == REG_DWORD || dwDataType == REG_DWORD_LITTLE_ENDIAN) {
			// Return a string representation of the decimal value.
			wchar_t sBinary[12];
			_ltow(*reinterpret_cast<const DWORD *>(pData), sBinary, 10);
			CComBSTR sData(sBinary);
			*psData = sData.Detach();
		} else if(dwDataType == REG_DWORD_BIG_ENDIAN)
			// Unsupported type.
			hr = E_ABORT;
		else if(dwDataType == REG_LINK)
			// Unsupported type.
			hr = E_ABORT;
		else if(dwDataType == REG_MULTI_SZ)
			// Unsupported type.
			hr = E_ABORT;
		else if(dwDataType == REG_NONE)
			// Unsupported type.
			hr = E_ABORT;
		else if(dwDataType == REG_QWORD || dwDataType == REG_QWORD_LITTLE_ENDIAN)
			// Unsupported type.
			hr = E_ABORT;

		if(hr == E_ABORT) {
			hr = RL_E_BADTYPEREAD;
			SetupErrorInfo(hr, (dwDataType == REG_DWORD_BIG_ENDIAN ? L"REG_DWORD_BIG_ENDIAN" : (dwDataType == REG_LINK ? L"REG_LINK" : (dwDataType == REG_MULTI_SZ ? L"REG_MULTI_SZ" : (dwDataType == REG_NONE ? L"REG_NONE" : (dwDataType == REG_QWORD ? L"REG_QWORD" : (dwDataType == REG_QWORD_LITTLE_ENDIAN ? L"REG_QWORD_LITTLE_ENDIAN" : L"UNKNOWN")))))));
		}
	} else {
		hr = RL_E_BADARG;
		SetupErrorInfo(hr, L"GenericRequestBroker::GetStringFromData");
	}

	return hr;
}
